import cv2
import numpy as np
import math
import subprocess
import os
import imageio_ffmpeg as ffmpeg


def is_cuda_available():
    """
    检查当前 OpenCV 是否支持 CUDA GPU
    :return: True 支持 CUDA, False 不支持
    """
    # OpenCV 是否编译了 CUDA
    try:
        count = cv2.cuda.getCudaEnabledDeviceCount()
        if count > 0:
            print(f"CUDA GPU 可用，设备数量: {count}")
            return True
        else:
            print("没有可用 CUDA GPU")
            return False
    except AttributeError:
        # cv2.cuda 不存在 → OpenCV 未编译 CUDA
        print("OpenCV 不支持 CUDA")
        return False


class VideoManager:
    def __init__(self, input_video, output_prefix="output_part_", split_duration=10,
                 watermark_area=None):
        """
        input_video: 输入视频路径
        output_prefix: 输出视频前缀
        split_duration: 每段切割时长（秒）
        watermark_area: 水印区域 (x1, y1, x2, y2)
        """
        self.input_video = input_video
        self.output_prefix = output_prefix
        self.split_duration = split_duration

        # 打开视频获取信息
        self.cap = cv2.VideoCapture(input_video)
        self.fps = self.cap.get(cv2.CAP_PROP_FPS)
        self.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        self.frame_count = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
        self.duration = self.frame_count / self.fps

        # 设置水印区域
        if watermark_area:
            self.logo_x1, self.logo_y1, self.logo_x2, self.logo_y2 = watermark_area
            self.logo_width = self.logo_x2 - self.logo_x1
            self.logo_height = self.logo_y2 - self.logo_y1
        else:
            # 默认右下角 200x100
            self.logo_x1, self.logo_y1 = self.width - 200, self.height - 100
            self.logo_x2, self.logo_y2 = self.width, self.height

        # 创建 mask
        self.mask = np.zeros((self.height, self.width), dtype=np.uint8)
        self.mask[self.logo_y1:self.logo_y2, self.logo_x1:self.logo_x2] = 255

        # 切割段数
        self.parts = math.ceil(self.duration / self.split_duration)

        # FFmpeg 路径
        self.ffmpeg_exe = ffmpeg.get_ffmpeg_exe()
        # self.ffmpeg_exe = "D:/ffmpeg-master-latest-win64-gpl/bin/ffmpeg.exe"
        print("使用 FFmpeg:", self.ffmpeg_exe)
        self.remove_watermark_path = f"video_remove_watermark.mp4"

    def save_first_frame(self, filename="first_frame.jpg"):
        self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
        ret, frame = self.cap.read()
        if ret:
            cv2.imwrite(filename, frame)
            print(f"第一帧已保存为: {filename}")
        else:
            print("无法读取视频第一帧")

    def remove_watermark_ffmpeg(self):
        """
        使用 FFmpeg delogo 滤镜去除视频水印
        :param use_gpu: 是否启用 GPU 加速（CUDA）
        :return: 去水印后的视频路径
        """
        cmd = [self.ffmpeg_exe, "-y"]  # 覆盖输出文件

        # 如果启用 GPU，加上 hwaccel 参数
        if is_cuda_available():
            cmd += ["-hwaccel", "cuda"]

        # 输入视频
        cmd += ["-i", self.input_video]

        # delogo 滤镜
        cmd += [
            "-vf",
            f"delogo=x={self.logo_x1}:y={self.logo_y1}:w={self.logo_width}:h={self.logo_height}:show=0",
            "-c:a", "copy",  # 保留原音频
            self.remove_watermark_path
        ]

        print("执行命令：", " ".join(cmd))
        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

        if result.returncode == 0:
            print(f"水印去除完成，输出文件：{self.remove_watermark_path}")
        else:
            print("去水印失败，错误信息：", result.stderr.decode())

    def remove_watermark(self, show_video=False):
        """
        视频水印去除
        :return:
        """
        fourcc = cv2.VideoWriter_fourcc(*"XVID")  # 用 AVI + XVID 避免黑屏
        writer = cv2.VideoWriter(self.remove_watermark_path, fourcc, self.fps,
                                 (self.width, self.height))
        x1, y1, x2, y2 = self.logo_x1, self.logo_y1, self.logo_x2, self.logo_y2
        roi_mask = self.mask[y1:y2, x1:x2]

        while True:
            ret, frame = self.cap.read()
            if not ret:
                break
            roi = frame[y1:y2, x1:x2]
            if is_cuda_available():
                # GPU 版本：上传到 GpuMat（OpenCV CUDA）
                gpu_roi = cv2.cuda_GpuMat()
                gpu_roi.upload(roi)
                # GPU 没有 inpaint，目前仍用 CPU 修复
                repaired = cv2.inpaint(roi, roi_mask, 3, cv2.INPAINT_TELEA)
                # dst = cv2.inpaint(src, inpaintMask, inpaintRadius, flags)
                #   src: 原始图像（彩色或灰度）
                #   inpaintMask: 掩码（同 src 大小，单通道，非零像素表示要修复的区域）
                #   inpaintRadius: 修复半径（3~5 常用，越大修复范围越广，但可能模糊）
                #   flags: 修复算法类型
                #       cv2.INPAINT_TELEA → 基于 Telea 算法，速度快，效果自然
                #       cv2.INPAINT_NS → 基于 Navier-Stokes 方程，效果更平滑但慢一些
                #   返回值 dst 是修复后的图像，即去掉水印后的帧。
            else:
                repaired = cv2.inpaint(roi, roi_mask, 3, cv2.INPAINT_TELEA)

            # 放回原帧
            frame[y1:y2, x1:x2] = repaired

            # 写入视频
            writer.write(frame)

            # ---------------------------
            # 2. 显示预览（可选）
            # ---------------------------
            if show_video:
                orig = frame.copy()
                cv2.rectangle(orig, (x1, y1), (x2, y2), (0, 255, 0), 2)
                combined = np.hstack((orig, frame))
                cv2.imshow("Original | Cleaned", combined)
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break

        writer.release()
        if show_video:
            cv2.destroyAllWindows()
        self.cap.release()

        print(f"水印去除完成: {self.remove_watermark_path}")

    def process(self):
        """
        1.视频去掉水印
        2.视频按时长切割成多个视频实现
        :return:
        """
        # ---------------------------
        # 1. Step 1：去水印
        self.remove_watermark_ffmpeg()  # 返回去水印后的视频路径
        # self.remove_watermark(show_video=True)
        print(f"[1/2] 去水印完成：{self.remove_watermark_path}")

        # ---------------------------
        # 2. Step 2：计算切割数量
        total_parts = int((self.duration + self.split_duration - 1) // self.split_duration)
        print(f"[2/2] 视频时长 {self.duration}s，将被切割为 {total_parts} 个片段，每段 {self.split_duration}s")

        # ---------------------------
        # 3. Step 3：循环切割
        for part_index in range(total_parts):
            self.video_cut(part_index)

        print("=== 视频处理完成：水印移除 + 切割完成 ===")
        if os.path.exists(self.remove_watermark_path):
            os.remove(self.remove_watermark_path)

    def _merge_audio(self, temp_video, part_index):
        output_name = f"{self.output_prefix}{part_index}.mp4"
        start_time = part_index * self.split_duration
        duration = min(self.split_duration, self.duration - start_time)
        print(f"合成音频: {output_name}, start={start_time}, duration={duration}")
        cmd = [
            self.ffmpeg_exe,  # FFmpeg 可执行文件路径
            "-ss", str(start_time),  # 从音频源的 start_time 秒开始
            "-y",  # 如果输出文件已存在，自动覆盖
            "-i", temp_video,  # 输入视频流（无音频的视频）
            "-i", self.input_video,  # 输入音频源（原视频）
            "-t", str(duration),  # 截取长度 duration 秒
            "-map", "0:v",  # 使用第一个输入（temp_video）的 **视频流**
            "-map", "1:a",  # 使用第二个输入（原视频）的 **音频流**
            "-c:v", "libx264",  # 视频流直接拷贝，不重新编码
            "-preset", "fast",
            "-crf", "23",
            "-c:a", "aac",  # 音频流编码为 AAC（保证兼容性）
            output_name  # 输出文件名
        ]
        subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    def video_cut(self, part_index):
        """
        视频切割
        :param part_index:
        :return:
        """
        output_name = f"{self.output_prefix}{part_index}.mp4"
        start_time = part_index * self.split_duration
        duration = min(self.split_duration, self.duration - start_time)
        print(f"切割视频: {output_name}, start={start_time}, duration={duration}")
        video_path = os.path.abspath(self.remove_watermark_path)
        cmd = [
            self.ffmpeg_exe,  # FFmpeg 可执行文件路径
            "-ss", str(start_time),  # 从音频源的 start_time 秒开始
            "-i", video_path,  # 输入音频源（原视频）
            "-t", str(duration),  # 截取长度 duration 秒
            "-c:v", "copy",  # 视频流直接拷贝，不重新编码
            output_name  # 输出文件名
        ]
        print("切割执行命令：", " ".join(cmd))
        subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)


# ---------------- 使用示例 ----------------
if __name__ == "__main__":
    video_manager = VideoManager(
        input_video="1.mp4",
        output_prefix="output_part_",
        split_duration=10 * 60,
        watermark_area=(555, 8, 618, 30)  # 可手动指定 (x1, y1, x2, y2)
    )
    # video_manager.save_first_frame("first_frame.jpg")  # 保存第一帧
    video_manager.process()  # 执行去水印 + 切割 + 水印标记
